# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved

from setuptools import setup, find_packages
from torch.utils import cpp_extension
import glob
from pathlib import Path
import os

root = Path(__file__).parent

sources = [
    "csrc/*.cpp",
    "csrc/utils/*.cpp",
    "csrc/aten/*.cpp",
    "csrc/comm/*.cpp",
    "csrc/comm/local/*.cpp",
    "csrc/tpp/alphafold/*.cpp",
]

sources = [j for i in sources for j in glob.glob(i)]
extra_compile_args = []
include_dirs = [(root / "csrc").as_posix()]

extra_compile_args += ["-fopenmp"]
if os.environ["KPEX_BUILD_TYPE"] == "release":
    extra_compile_args += ["-O3"]
elif os.environ["KPEX_BUILD_TYPE"] == "debug":
    extra_compile_args += ["-O0", "-g"]
else:
    print("requires env KPEX_BUILD_TYPE (release / debug)")

library_dirs = []
libraries = []
KUTACC_ROOT = os.environ.get("KUTACC_ROOT", None)
if KUTACC_ROOT:
    include_dirs += [f"{KUTACC_ROOT}/include"]
    library_dirs = [f"{KUTACC_ROOT}/lib"]
    libraries = ["kutacc"]
    extra_compile_args += [f"-Wl, -rpath={KUTACC_ROOT}/lib"]

setup(
    name="kunpeng-pytorch-extension",
    version="0.0.1",
    author="KPEX",
    packages=find_packages(),
    ext_modules=[
        cpp_extension.CppExtension(
            name="kpex._C",
            sources=sources,
            include_dirs=include_dirs,
            libraries=libraries,
            library_dirs=library_dirs,
            extra_compile_args={"cxx": extra_compile_args},
        )
    ],
    cmdclass={"build_ext": cpp_extension.BuildExtension}
)
